"""
Alpha051-075因子使用示例
=======================

展示如何使用Alpha101因子库中第三批因子Alpha051-075的完整示例。
包括数据准备、因子计算、结果分析和可视化。
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime, timedelta
import warnings
import sys
import os

# 添加项目路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from alpha101.core import Alpha101Factory, AlphaFactor
from alpha101.config import get_factor_config, list_factors_by_category

# 设置中文字体和样式
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False
sns.set_style("whitegrid")
warnings.filterwarnings('ignore')

def generate_realistic_data(n_stocks=30, n_days=252):
    """
    生成更真实的股票数据
    
    Parameters:
    -----------
    n_stocks : int
        股票数量
    n_days : int
        交易日数量
        
    Returns:
    --------
    dict
        股票数据字典
    pd.DatetimeIndex
        日期索引
    list
        股票代码列表
    """
    np.random.seed(42)
    
    # 生成日期索引
    start_date = datetime(2023, 1, 1)
    dates = pd.date_range(start_date, periods=n_days, freq='B')  # 工作日
    
    # 生成股票代码
    stock_codes = [f"Stock_{i:03d}" for i in range(1, n_stocks + 1)]
    
    # 生成基础收益率（带有市场因子）
    market_returns = np.random.normal(0.0005, 0.015, n_days)
    
    # 生成个股特异收益率
    idiosyncratic_returns = np.random.normal(0, 0.02, (n_stocks, n_days))
    
    # 生成beta系数
    betas = np.random.uniform(0.5, 1.5, n_stocks)
    
    # 计算总收益率
    returns = np.zeros((n_stocks, n_days))
    for i in range(n_stocks):
        returns[i, :] = betas[i] * market_returns + idiosyncratic_returns[i, :]
    
    # 生成价格序列
    base_prices = np.random.uniform(10, 100, n_stocks)
    prices = np.zeros((n_stocks, n_days))
    prices[:, 0] = base_prices
    
    for t in range(1, n_days):
        prices[:, t] = prices[:, t-1] * (1 + returns[:, t])
    
    # 生成OHLC数据
    close = prices
    
    # 生成开盘价（基于前一日收盘价加上隔夜收益）
    overnight_returns = np.random.normal(0, 0.005, (n_stocks, n_days))
    open_price = np.zeros_like(close)
    open_price[:, 0] = close[:, 0] * (1 + overnight_returns[:, 0])
    for t in range(1, n_days):
        open_price[:, t] = close[:, t-1] * (1 + overnight_returns[:, t])
    
    # 生成日内波动
    intraday_volatility = np.random.uniform(0.01, 0.03, (n_stocks, n_days))
    high = np.maximum(open_price, close) * (1 + intraday_volatility)
    low = np.minimum(open_price, close) * (1 - intraday_volatility)
    
    # 生成成交量（与价格变化相关）
    volume_base = np.random.lognormal(12, 1, n_stocks)
    volume = np.zeros((n_stocks, n_days))
    for i in range(n_stocks):
        for t in range(n_days):
            # 成交量与价格变化幅度正相关
            price_change = abs(returns[i, t])
            volume_multiplier = 1 + price_change * 10
            volume[i, t] = volume_base[i] * volume_multiplier * np.random.lognormal(0, 0.3)
    
    # 生成VWAP（成交量加权平均价格）
    vwap = (high + low + close * 2) / 4  # 简化的VWAP计算
    
    return {
        'open': open_price,
        'high': high,
        'low': low,
        'close': close,
        'volume': volume,
        'vwap': vwap
    }, dates, stock_codes

def calculate_alpha051_075_factors(data):
    """
    计算Alpha051-075所有因子
    
    Parameters:
    -----------
    data : dict
        股票数据
        
    Returns:
    --------
    dict
        因子计算结果
    """
    print("开始计算Alpha051-075因子...")
    
    # 创建因子工厂
    factory = Alpha101Factory(use_numba=False)
    
    # Alpha051-075因子列表
    factor_ids = [f"alpha{i:03d}" for i in range(51, 76)]
    
    results = {}
    successful_factors = []
    failed_factors = []
    
    for i, factor_id in enumerate(factor_ids, 1):
        try:
            print(f"计算 {factor_id} ({i}/25)...")
            
            # 计算因子
            factor_values = factory.compute_factor(
                factor_id, 
                data,
                winsorize=True,
                standardize=True
            )
            
            if factor_values is not None:
                results[factor_id] = factor_values
                successful_factors.append(factor_id)
                print(f"  ✓ {factor_id} 计算成功")
            else:
                failed_factors.append(factor_id)
                print(f"  ✗ {factor_id} 计算失败")
                
        except Exception as e:
            failed_factors.append(factor_id)
            print(f"  ✗ {factor_id} 计算失败: {str(e)}")
    
    print(f"\n计算完成: {len(successful_factors)}/25 个因子成功")
    if failed_factors:
        print(f"失败的因子: {', '.join(failed_factors)}")
    
    return results, successful_factors, failed_factors

def analyze_factor_properties(results, dates, stock_codes):
    """
    分析因子属性
    
    Parameters:
    -----------
    results : dict
        因子计算结果
    dates : pd.DatetimeIndex
        日期索引
    stock_codes : list
        股票代码列表
        
    Returns:
    --------
    pd.DataFrame
        因子属性分析结果
    """
    print("\n分析因子属性...")
    
    analysis_results = []
    
    for factor_id, factor_values in results.items():
        try:
            config = get_factor_config(factor_id)
            
            # 计算基本统计量
            valid_values = factor_values[~np.isnan(factor_values)]
            
            if len(valid_values) > 0:
                analysis = {
                    'factor_id': factor_id,
                    'description': config['description'],
                    'category': config['category'],
                    'window': config['window'],
                    'valid_ratio': len(valid_values) / factor_values.size,
                    'mean': np.mean(valid_values),
                    'std': np.std(valid_values),
                    'skewness': pd.Series(valid_values).skew(),
                    'kurtosis': pd.Series(valid_values).kurtosis(),
                    'min_value': np.min(valid_values),
                    'max_value': np.max(valid_values),
                    'q25': np.percentile(valid_values, 25),
                    'q50': np.percentile(valid_values, 50),
                    'q75': np.percentile(valid_values, 75)
                }
                analysis_results.append(analysis)
        except Exception as e:
            print(f"分析 {factor_id} 时出错: {str(e)}")
    
    return pd.DataFrame(analysis_results)

def create_factor_correlation_matrix(results):
    """
    创建因子相关性矩阵
    
    Parameters:
    -----------
    results : dict
        因子计算结果
        
    Returns:
    --------
    pd.DataFrame
        相关性矩阵
    """
    print("计算因子相关性矩阵...")
    
    # 将因子数据转换为DataFrame
    factor_data = {}
    for factor_id, factor_values in results.items():
        # 取最后一天的截面数据
        last_day_values = factor_values[:, -1]
        valid_mask = ~np.isnan(last_day_values)
        if np.sum(valid_mask) > 10:  # 至少需要10个有效值
            factor_data[factor_id] = last_day_values
    
    if len(factor_data) > 1:
        df = pd.DataFrame(factor_data)
        correlation_matrix = df.corr()
        return correlation_matrix
    else:
        return pd.DataFrame()

def visualize_results(results, analysis_df, correlation_matrix, dates, stock_codes):
    """
    可视化结果
    
    Parameters:
    -----------
    results : dict
        因子计算结果
    analysis_df : pd.DataFrame
        因子属性分析结果
    correlation_matrix : pd.DataFrame
        相关性矩阵
    dates : pd.DatetimeIndex
        日期索引
    stock_codes : list
        股票代码列表
    """
    print("创建可视化图表...")
    
    # 设置图形大小
    plt.style.use('default')
    
    # 1. 因子分类分布
    if not analysis_df.empty:
        plt.figure(figsize=(12, 8))
        
        # 子图1: 因子分类分布
        plt.subplot(2, 3, 1)
        category_counts = analysis_df['category'].value_counts()
        plt.pie(category_counts.values, labels=category_counts.index, autopct='%1.1f%%')
        plt.title('Alpha051-075因子分类分布')
        
        # 子图2: 有效值比例分布
        plt.subplot(2, 3, 2)
        plt.hist(analysis_df['valid_ratio'], bins=10, alpha=0.7, color='skyblue')
        plt.xlabel('有效值比例')
        plt.ylabel('因子数量')
        plt.title('因子有效值比例分布')
        
        # 子图3: 因子标准差分布
        plt.subplot(2, 3, 3)
        plt.hist(analysis_df['std'], bins=10, alpha=0.7, color='lightgreen')
        plt.xlabel('标准差')
        plt.ylabel('因子数量')
        plt.title('因子标准差分布')
        
        # 子图4: 偏度分布
        plt.subplot(2, 3, 4)
        plt.hist(analysis_df['skewness'].dropna(), bins=10, alpha=0.7, color='orange')
        plt.xlabel('偏度')
        plt.ylabel('因子数量')
        plt.title('因子偏度分布')
        
        # 子图5: 峰度分布
        plt.subplot(2, 3, 5)
        plt.hist(analysis_df['kurtosis'].dropna(), bins=10, alpha=0.7, color='pink')
        plt.xlabel('峰度')
        plt.ylabel('因子数量')
        plt.title('因子峰度分布')
        
        # 子图6: 时间窗口分布
        plt.subplot(2, 3, 6)
        window_counts = analysis_df['window'].value_counts().sort_index()
        plt.bar(window_counts.index, window_counts.values, alpha=0.7, color='purple')
        plt.xlabel('时间窗口')
        plt.ylabel('因子数量')
        plt.title('因子时间窗口分布')
        
        plt.tight_layout()
        plt.savefig('factor/examples/alpha051_075_analysis.png', dpi=300, bbox_inches='tight')
        plt.show()
    
    # 2. 相关性热力图
    if not correlation_matrix.empty and correlation_matrix.shape[0] > 1:
        plt.figure(figsize=(12, 10))
        mask = np.triu(np.ones_like(correlation_matrix, dtype=bool))
        sns.heatmap(correlation_matrix, mask=mask, annot=True, cmap='coolwarm', 
                   center=0, square=True, fmt='.2f', cbar_kws={"shrink": .8})
        plt.title('Alpha051-075因子相关性矩阵')
        plt.tight_layout()
        plt.savefig('factor/examples/alpha051_075_correlation.png', dpi=300, bbox_inches='tight')
        plt.show()
    
    # 3. 展示几个代表性因子的时间序列
    if results:
        plt.figure(figsize=(15, 10))
        
        # 选择几个代表性因子
        sample_factors = list(results.keys())[:6]  # 取前6个成功的因子
        
        for i, factor_id in enumerate(sample_factors, 1):
            plt.subplot(2, 3, i)
            factor_values = results[factor_id]
            
            # 计算每日因子均值
            daily_means = np.nanmean(factor_values, axis=0)
            valid_dates = dates[:len(daily_means)]
            
            plt.plot(valid_dates, daily_means, linewidth=1.5)
            plt.title(f'{factor_id}')
            plt.xlabel('日期')
            plt.ylabel('因子值')
            plt.xticks(rotation=45)
            
            # 添加零线
            plt.axhline(y=0, color='red', linestyle='--', alpha=0.5)
        
        plt.tight_layout()
        plt.savefig('factor/examples/alpha051_075_timeseries.png', dpi=300, bbox_inches='tight')
        plt.show()

def demonstrate_single_factor_usage():
    """
    演示单个因子的详细使用方法
    """
    print("\n=== 单个因子使用演示 ===")
    
    # 生成测试数据
    data, dates, stock_codes = generate_realistic_data(n_stocks=20, n_days=100)
    
    # 创建因子工厂
    factory = Alpha101Factory(use_numba=False)
    
    # 演示Alpha051因子
    factor_id = "alpha051"
    print(f"\n演示因子: {factor_id}")
    
    # 获取因子信息
    config = get_factor_config(factor_id)
    print(f"因子描述: {config['description']}")
    print(f"因子公式: {config['formula']}")
    print(f"数据依赖: {config['dependencies']}")
    print(f"时间窗口: {config['window']}")
    print(f"因子类别: {config['category']}")
    
    # 计算因子
    try:
        factor_values = factory.compute_factor(factor_id, data)
        
        if factor_values is not None:
            print(f"\n因子计算成功!")
            print(f"结果形状: {factor_values.shape}")
            print(f"有效值数量: {np.sum(~np.isnan(factor_values))}")
            print(f"有效值比例: {np.sum(~np.isnan(factor_values)) / factor_values.size:.2%}")
            
            # 显示最后一天的因子值
            last_day_values = factor_values[:, -1]
            valid_mask = ~np.isnan(last_day_values)
            
            if np.sum(valid_mask) > 0:
                print(f"\n最后一天的因子值 (前10个股票):")
                for i in range(min(10, len(stock_codes))):
                    if valid_mask[i]:
                        print(f"  {stock_codes[i]}: {last_day_values[i]:.6f}")
                    else:
                        print(f"  {stock_codes[i]}: NaN")
        else:
            print("因子计算失败!")
            
    except Exception as e:
        print(f"因子计算出错: {str(e)}")

def main():
    """主函数"""
    print("Alpha051-075因子使用示例")
    print("=" * 50)
    
    # 1. 生成测试数据
    print("1. 生成测试数据...")
    data, dates, stock_codes = generate_realistic_data(n_stocks=30, n_days=252)
    print(f"   数据形状: {data['close'].shape}")
    print(f"   日期范围: {dates[0].strftime('%Y-%m-%d')} 到 {dates[-1].strftime('%Y-%m-%d')}")
    print(f"   股票数量: {len(stock_codes)}")
    
    # 2. 计算所有Alpha051-075因子
    print("\n2. 计算Alpha051-075因子...")
    results, successful_factors, failed_factors = calculate_alpha051_075_factors(data)
    
    if not results:
        print("没有成功计算的因子，程序退出。")
        return
    
    # 3. 分析因子属性
    print("\n3. 分析因子属性...")
    analysis_df = analyze_factor_properties(results, dates, stock_codes)
    
    if not analysis_df.empty:
        print("\n因子属性统计:")
        print(analysis_df[['factor_id', 'category', 'valid_ratio', 'mean', 'std']].round(4))
        
        # 按类别统计
        print("\n按类别统计:")
        category_stats = analysis_df.groupby('category').agg({
            'factor_id': 'count',
            'valid_ratio': 'mean',
            'std': 'mean'
        }).round(4)
        category_stats.columns = ['因子数量', '平均有效值比例', '平均标准差']
        print(category_stats)
    
    # 4. 计算相关性矩阵
    print("\n4. 计算因子相关性...")
    correlation_matrix = create_factor_correlation_matrix(results)
    
    if not correlation_matrix.empty:
        print("相关性矩阵计算完成")
        print(f"矩阵大小: {correlation_matrix.shape}")
        
        # 显示高相关性因子对
        high_corr_pairs = []
        for i in range(len(correlation_matrix.columns)):
            for j in range(i+1, len(correlation_matrix.columns)):
                corr_val = correlation_matrix.iloc[i, j]
                if abs(corr_val) > 0.7:  # 高相关性阈值
                    high_corr_pairs.append((
                        correlation_matrix.columns[i],
                        correlation_matrix.columns[j],
                        corr_val
                    ))
        
        if high_corr_pairs:
            print(f"\n发现 {len(high_corr_pairs)} 对高相关性因子 (|相关系数| > 0.7):")
            for factor1, factor2, corr in high_corr_pairs:
                print(f"  {factor1} - {factor2}: {corr:.3f}")
    
    # 5. 可视化结果
    print("\n5. 创建可视化图表...")
    try:
        visualize_results(results, analysis_df, correlation_matrix, dates, stock_codes)
        print("可视化图表已保存到 factor/examples/ 目录")
    except Exception as e:
        print(f"可视化过程中出错: {str(e)}")
    
    # 6. 演示单个因子使用
    demonstrate_single_factor_usage()
    
    print("\n" + "=" * 50)
    print("Alpha051-075因子使用示例完成!")
    print(f"成功计算 {len(successful_factors)} 个因子")
    print("详细结果已保存为图片文件")

if __name__ == "__main__":
    main()